cc @owenowenisme

Hi @400Ping, thanks for your contribution! We really appreciate the effort.

I think @owenowenisme might be right here, all_inputs_done should already be called. My hunch is that BlockRefBundler class is not correctly finalizing the outputs. From here

            if (
                output_buffer_size < self._min_rows_per_bundle
                or output_buffer_size == 0
            ):

it looks like we should be doing

            if (
                output_buffer_size < self._min_rows_per_bundle
                or output_buffer_size == 0
                or self._finalized
            ):

so that there are no more remainders. Can you try that and report back?
cc: @bveeramani

Hey @400Ping, would you mind helping me understand the root cause of the assertion error?

Also, did the repro script fail before the changes?

Hey @400Ping, would you mind helping me understand the root cause of the assertion error?

Also, did the repro script fail before the changes?

Ok, will try to find it.

cc @bveeramani PTAL

not very sure if this is the way to solve this.

@400Ping @iamjustinhsu I was able to create a minimal repro of this issue (at least, I think this is the same issue)

import ray


class Fn:

    def __call__(self, batch):
        return batch


ds = ray.data.range(100, override_num_blocks=100).map_batches(Fn, batch_size=10).limit(1)
for _ in ds.iter_internal_ref_bundles():
    pass

Traceback (most recent call last):                                                                                                                                      
  File "/Users/balaji/ray/1.py", line 12, in <module>
    for _ in ds.iter_internal_ref_bundles():                                                                                                                            
  File "/Users/balaji/ray/python/ray/data/_internal/execution/interfaces/executor.py", line 34, in __next__
    return self.get_next()
  File "/Users/balaji/ray/python/ray/data/_internal/execution/legacy_compat.py", line 76, in get_next
    bundle = self._base_iterator.get_next(output_split_idx)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 786, in get_next
    bundle = state.get_output_blocking(output_split_idx)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor_state.py", line 454, in get_output_blocking
    raise self._exception
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 356, in run
    continue_sched = self._scheduling_loop_step(self._topology)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 529, in _scheduling_loop_step
    self._validate_operator_queues_empty(op, state)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 571, in _validate_operator_queues_empty
    assert op.internal_input_queue_num_blocks() == 0, error_msg.format(
AssertionError: Expected Internal Input Queue for MapBatches(Fn) to be empty, but found 8 bundles

@iamjustinhsu would a reasonable fix be to just clear the internal queues when the operator is manually marked finished?

    def mark_execution_finished(self):
        # Discard remaining bundles in the internal bundle queue.
        self._bundle_queue.clear()

        # Discard remaining bundles in the block ref bundler.
        self._block_ref_bundler.done_adding_bundles()
        while self._block_ref_bundler.has_bundle():
            self._block_ref_bundler.get_next_bundle()

        super().mark_execution_finished()

@bveeramani oh i see now. We actually do this for for all InternalQueueOperatorMixin operators. However, I implemented this for MapOperator, thinking ActorPoolMapOperator uses the same underlying implementation. It actually doesn't, and contains an additional queue called _bundle_queue. Here is my PR for reference: #58441. What we should do is have logic for also clearing the self._bundle_queue in ActorPoolMapOperator. This can be done with having separate implementations for TaskPoolMapOperator and ActorPoolMapOperator

Ah, okay.

@400Ping would you mind refactoring the PR or opening a new PR to do the following:

• Delete MapOperator.clear_internal_output_queue and MapOperator.clear_internal_input_queue
• Add TaskPoolMapOperator.clear_internal_output_queue and TaskPoolMapOperator.clear_internal_input_queue
• Add ActorPoolMapOperator.clear_internal_output_queue and ActorPoolMapOperator.clear_internal_input_queue

I think this is all that we need to do to solve the flaky issue.

The implementations for clear_internal_input_queue and clear_internal_output_queue LGTM.

@400Ping to keep the git history clear, would you mind reverting all of the changes unrelated to those methods (e.g., formatting, removing _inputs_done, type annotations)? Once we've reverted the unrelated changes, I'll approve the PR

Ok.

Thanks for the fix!

@400Ping Just merged! ty for the contribution!

@400Ping Just merged! ty for the contribution!

Thank you as well, I am a newbie in this area 😓.